"""
    EnKFDiagnostics

Lightweight diagnostic helpers used by the reduced EnKF driver. The utilities
here intentionally avoid heavy dependencies so that unit tests can exercise the
analysis workflow without requiring the full operational plotting and logging
stack.
"""
module EnKFDiagnostics

using Dates
using Statistics

export DiagnosticConfig, InnovationStats, EnsembleSpreadStats, PerformanceStats
export MemoryStats, print_innovation_stats, calculate_ensemble_spread
export monitor_performance, track_memory_usage, validate_ensemble_consistency
export generate_diagnostic_report

struct DiagnosticConfig
    output_path::String
    log_level::Int
    save_statistics::Bool
    generate_plots::Bool
    observation_types::Vector{String}
    variable_names::Vector{String}
    regional_analysis::Bool
    latitude_bands::Vector{Float64}
    quality_control_thresholds::Dict{String,Float64}
    performance_monitoring::Bool
    memory_tracking::Bool
end

function DiagnosticConfig(; output_path::AbstractString,
                          log_level::Integer,
                          save_statistics::Bool,
                          generate_plots::Bool,
                          observation_types::AbstractVector{<:AbstractString},
                          variable_names::AbstractVector{<:AbstractString},
                          regional_analysis::Bool,
                          latitude_bands::AbstractVector{<:Real},
                          quality_control_thresholds::Dict{String,Float64},
                          performance_monitoring::Bool,
                          memory_tracking::Bool)
    return DiagnosticConfig(
        String(output_path),
        Int(log_level),
        save_statistics,
        generate_plots,
        String.(observation_types),
        String.(variable_names),
        regional_analysis,
        Float64.(latitude_bands),
        quality_control_thresholds,
        performance_monitoring,
        memory_tracking,
    )
end

mutable struct InnovationStats
    label::String
    total_observations::Int
    prior_mean::Float64
    prior_std::Float64
    prior_rms::Float64
    prior_bias::Float64
end

InnovationStats(label::String) = InnovationStats(label, 0, 0.0, 0.0, 0.0, 0.0)

struct EnsembleSpreadStats
    variable::String
    mean_spread::Float64
    min_spread::Float64
    max_spread::Float64
    spread_std::Float64
    vertical_spread_profile::Vector{Float64}
end

struct PerformanceStats
    total_runtime::Float64
    data_reading_time::Float64
    ensemble_update_time::Float64
end

struct MemoryStats
    peak_memory_usage::Float64
    current_memory_usage::Float64
end

function print_innovation_stats(innovation_fits::Vector{Float64},
                                ensemble_spread::Vector{Float64},
                                observation_errors::Vector{Float64},
                                config::DiagnosticConfig)
    println("\n" * "="^60)
    println("ENSEMBLE INNOVATION STATISTICS (compact)")
    println("="^60)

    if isempty(innovation_fits)
        println("No observations supplied — skipping statistics")
        return InnovationStats("empty")
    end

    innovation_mean = mean(innovation_fits)
    innovation_std = std(innovation_fits)
    innovation_rms = sqrt(mean(innovation_fits .^ 2))
    obs_error_mean = isempty(observation_errors) ? 0.0 : mean(observation_errors)
    spread_mean = isempty(ensemble_spread) ? 0.0 : mean(ensemble_spread)

    println("Count: $(length(innovation_fits))")
    println("Innovation mean: $(round(innovation_mean, digits=4))")
    println("Innovation RMS:  $(round(innovation_rms, digits=4))")
    println("Spread mean:     $(round(spread_mean, digits=4))")
    println("Obs error mean:  $(round(obs_error_mean, digits=4))")

    stats = InnovationStats("combined")
    stats.total_observations = length(innovation_fits)
    stats.prior_mean = innovation_mean
    stats.prior_std = innovation_std
    stats.prior_rms = innovation_rms
    stats.prior_bias = innovation_mean
    return stats
end

function calculate_ensemble_spread(ensemble_data::Array{<:Real,4}, variable_names::Vector{String})
    nlons, nlats, nlevs, nmembers = size(ensemble_data)
    results = Dict{String,EnsembleSpreadStats}()

    spread = std(ensemble_data; dims = 4, corrected = true)[:,:,:,1]
    mean_spread = mean(spread)
    min_spread = minimum(spread)
    max_spread = maximum(spread)
    spread_std = std(spread)
    profile = [mean(view(spread, :, :, k)) for k in 1:nlevs]

    var_name = isempty(variable_names) ? "state" : variable_names[1]
    stats = EnsembleSpreadStats(var_name, mean_spread, min_spread, max_spread, spread_std, profile)
    results[var_name] = stats
    return results
end

function monitor_performance(start_time::DateTime, stage::AbstractString, config::DiagnosticConfig)
    elapsed = (now() - start_time).value / 1_000
    return PerformanceStats(elapsed, elapsed, elapsed)
end

function track_memory_usage(config::DiagnosticConfig)
    bytes = Base.gc_bytes()
    mb = bytes / 1_048_576
    return MemoryStats(mb, mb)
end

function validate_ensemble_consistency(ensemble_data::Array{<:Real,4}, config::DiagnosticConfig)
    has_nan = any(isnan, ensemble_data)
    has_inf = any(isinf, ensemble_data)
    return Dict(
        "no_invalid_values" => !(has_nan || has_inf),
        "member_count" => size(ensemble_data, 4)
    )
end

function generate_diagnostic_report(all_stats::Dict, config::DiagnosticConfig)
    report_path = joinpath(config.output_path, "enkf_diagnostic_report.txt")
    open(report_path, "w") do io
        println(io, "EnKF Diagnostic Report (minimal)")
        println(io, "Generated: $(Dates.format(now(), "yyyy-mm-dd HH:MM:SS"))")
        println(io)

        if haskey(all_stats, "innovation")
            println(io, "Innovation summary:")
            for (label, stats) in all_stats["innovation"]
                println(io, "  $label: mean=$(stats.prior_mean), rms=$(stats.prior_rms), N=$(stats.total_observations)")
            end
            println(io)
        end

        if haskey(all_stats, "ensemble_spread")
            println(io, "Ensemble spread summary:")
            for (label, stats) in all_stats["ensemble_spread"]
                println(io, "  $label: mean=$(stats.mean_spread), min=$(stats.min_spread), max=$(stats.max_spread)")
            end
            println(io)
        end
    end
    return report_path
end

end # module EnKFDiagnostics
